import json
import pandas as pd
import numpy as np
from datetime import datetime
from django.shortcuts import render
from django.db import transaction
from valar.classes.valar_response import ValarResponse
from valar.dao import Dao

from .models import (
    ForecastTask,
    TimeSeriesData,
    ForecastResult,
    HierarchicalStructure,
)

# 导入核心预测算法
from src.data_loader import DataLoader
from src.frameworks.framework1 import Framework1
from src.frameworks.framework2 import Framework2
from src.frameworks.framework3 import Framework3
from src.evaluation import Evaluation


def create_forecast_task(request):
    """创建新的预测任务"""
    try:
        body = json.loads(request.body)

        # 创建预测任务
        task = ForecastTask.objects.create(
            task_name=body.get('task_name', ''),
            description=body.get('description', ''),
            forecast_horizon=body.get('forecast_horizon', 5),
            start_year=body.get('start_year'),
            end_year=body.get('end_year'),
        )

        return ValarResponse(True, data={'task_id': task.id})

    except Exception as e:
        return ValarResponse(False, message=f"创建任务失败: {str(e)}")


def upload_data(request):
    """上传时间序列数据"""
    try:
        body = json.loads(request.body)
        task_id = body.get('task_id')
        data_records = body.get('data', [])

        if not task_id:
            return ValarResponse(False, message="缺少任务ID")

        try:
            task = ForecastTask.objects.get(id=task_id)
        except ForecastTask.DoesNotExist:
            return ValarResponse(False, message="任务不存在")

        # 批量创建数据记录
        data_objects = []
        for record in data_records:
            data_objects.append(TimeSeriesData(
                task=task,
                year=record['year'],
                province=record['province'],
                region=record.get('region', ''),
                value=record['value'],
                data_type='original'
            ))

        TimeSeriesData.objects.bulk_create(data_objects, ignore_conflicts=True)

        return ValarResponse(True, data={'uploaded_count': len(data_objects)})

    except Exception as e:
        return ValarResponse(False, message=f"数据上传失败: {str(e)}")


def run_bottom_up_forecast(request):
    """执行自底向上预测"""
    try:
        body = json.loads(request.body)
        task_id = body.get('task_id')

        if not task_id:
            return ValarResponse(False, message="缺少任务ID")

        try:
            task = ForecastTask.objects.get(id=task_id)
        except ForecastTask.DoesNotExist:
            return ValarResponse(False, message="任务不存在")

        # 更新任务状态
        task.status = 'running'
        task.save()

        try:
            # 1. 加载数据
            data_df = _load_task_data(task)

            # 2. 执行自底向上预测
            result = _execute_bottom_up_forecast(data_df, task.forecast_horizon)

            if result['success']:
                # 3. 保存预测结果
                _save_forecast_results(task, 'bottom_up', result['data'])

                # 4. 更新任务状态
                task.status = 'completed'
                task.save()
                return ValarResponse(True, data=result['data'])
            else:
                # 预测失败
                task.status = 'failed'
                task.save()

                return ValarResponse(False, message=result['error'])

        except Exception as e:
            task.status = 'failed'
            task.save()

            raise e

    except Exception as e:
        return ValarResponse(False, message=f"自底向上预测失败: {str(e)}")


def run_mint_forecast(request):
    """执行MinT协调预测"""
    try:
        body = json.loads(request.body)
        task_id = body.get('task_id')

        if not task_id:
            return ValarResponse(False, message="缺少任务ID")

        try:
            task = ForecastTask.objects.get(id=task_id)
        except ForecastTask.DoesNotExist:
            return ValarResponse(False, message="任务不存在")

        task.status = 'running'
        task.save()


        try:
            # 加载数据
            data_df = _load_task_data(task)

            # 执行MinT预测
            result = _execute_mint_forecast(data_df, task.forecast_horizon)

            if result['success']:
                _save_forecast_results(task, 'mint', result['data'])

                task.status = 'completed'
                task.save()


                return ValarResponse(True, data=result['data'])
            else:
                task.status = 'failed'
                task.save()

                return ValarResponse(False, message=result['error'])

        except Exception as e:
            task.status = 'failed'
            task.save()


            raise e

    except Exception as e:
        return ValarResponse(False, message=f"MinT预测失败: {str(e)}")


def run_ilr_var_forecast(request):
    """执行ILR-VAR预测"""
    try:
        body = json.loads(request.body)
        task_id = body.get('task_id')

        if not task_id:
            return ValarResponse(False, message="缺少任务ID")

        try:
            task = ForecastTask.objects.get(id=task_id)
        except ForecastTask.DoesNotExist:
            return ValarResponse(False, message="任务不存在")

        task.status = 'running'
        task.save()

        try:
            # 加载数据
            data_df = _load_task_data(task)

            # 执行ILR-VAR预测
            result = _execute_ilr_var_forecast(data_df, task.forecast_horizon)

            if result['success']:
                _save_forecast_results(task, 'ilr_var', result['data'])

                task.status = 'completed'
                task.save()


                return ValarResponse(True, data=result['data'])
            else:
                task.status = 'failed'
                task.save()


                return ValarResponse(False, message=result['error'])

        except Exception as e:
            task.status = 'failed'
            task.save()

            raise e

    except Exception as e:
        return ValarResponse(False, message=f"ILR-VAR预测失败: {str(e)}")


def get_forecast_results(request):
    """获取预测结果"""
    try:
        body = json.loads(request.body)
        task_id = body.get('task_id')
        framework = body.get('framework')  # 可选，指定框架类型

        if not task_id:
            return ValarResponse(False, message="缺少任务ID")

        try:
            task = ForecastTask.objects.get(id=task_id)
        except ForecastTask.DoesNotExist:
            return ValarResponse(False, message="任务不存在")

        # 构建查询条件
        query_params = {'task': task}
        if framework:
            query_params['framework'] = framework

        # 获取预测结果
        results = ForecastResult.objects.filter(**query_params).values(
            'framework', 'forecast_year', 'province', 'region',
            'forecast_value', 'lower_bound', 'upper_bound', 'created_date'
        )

        return ValarResponse(True, data=list(results))

    except Exception as e:
        return ValarResponse(False, message=f"获取结果失败: {str(e)}")



# 辅助函数
def _load_task_data(task):
    """加载任务的时间序列数据"""
    data = TimeSeriesData.objects.filter(
        task=task,
        data_type='original'
    ).values('year', 'province', 'value')

    df = pd.DataFrame(list(data))
    if df.empty:
        raise ValueError("任务数据为空")

    # 转换为宽格式
    df_pivot = df.pivot(index='year', columns='province', values='value')
    df_pivot.fillna(0, inplace=True)

    return df_pivot


def _execute_bottom_up_forecast(data_df, horizon):
    """执行自底向上预测"""
    try:
        from statsmodels.tsa.arima.model import ARIMA
        
        forecasts = {}
        forecast_lower = {}
        forecast_upper = {}

        # 对每个省份进行ARIMA预测
        for province in data_df.columns:
            series = data_df[province].dropna()
            series = series[series > 0]  # 只保留正值
            
            if len(series) < 10:
                print(f"跳过 {province}：数据点不足 ({len(series)} < 10)")
                continue

            try:
                # 使用简单的ARIMA(1,1,1)模型
                model = ARIMA(series, order=(1, 1, 1))
                fitted_model = model.fit()
                
                # 预测
                forecast_result = fitted_model.forecast(steps=horizon)
                
                # 简单的置信区间
                std_error = np.std(series) * 0.1
                lower_bound = forecast_result - 1.96 * std_error
                upper_bound = forecast_result + 1.96 * std_error
                
                forecasts[province] = forecast_result.values
                forecast_lower[province] = lower_bound.values
                forecast_upper[province] = upper_bound.values
                
            except Exception as e:
                print(f"预测 {province} 失败: {str(e)}")
                continue

        if not forecasts:
            return {'success': False, 'error': '没有成功预测任何省份'}

        # 构建结果数据
        result_data = []
        last_year = data_df.index.max()

        for i in range(horizon):
            forecast_year = last_year + i + 1
            for province in forecasts.keys():
                result_data.append({
                    'forecast_year': forecast_year,
                    'province': province,
                    'forecast_value': float(forecasts[province][i]),
                    'lower_bound': float(forecast_lower[province][i]),
                    'upper_bound': float(forecast_upper[province][i])
                })

        return {'success': True, 'data': result_data}

    except Exception as e:
        return {'success': False, 'error': str(e)}


def _execute_mint_forecast(data_df, horizon):
    """执行MinT预测"""
    try:
        # 先准备数据格式
        from src.data_loader import DataLoader
        
        # 创建数据加载器并处理数据
        data_loader = DataLoader()
        
        # 将DataFrame转换为DataLoader期望的格式
        processed_data = data_loader.process_dataframe(data_df)
        
        # 包装数据以符合Framework2的期望格式
        wrapped_data = {
            'Sheet1': processed_data
        }
        
        # 使用Framework2进行MinT预测
        framework2 = Framework2(wrapped_data)

        # 执行预测
        result = framework2.forecast(horizon)

        if result is not None:
            # 转换结果格式
            result_data = []
            last_year = data_df.index.max()

            for i in range(horizon):
                forecast_year = last_year + i + 1
                for province in result.columns:
                    if not province.endswith(('_yhat_lower', '_yhat_upper')):
                        forecast_value = float(result.loc[result.index[i], province])
                        lower_col = f'{province}_yhat_lower'
                        upper_col = f'{province}_yhat_upper'

                        lower_bound = float(
                            result.loc[result.index[i], lower_col]) if lower_col in result.columns else None
                        upper_bound = float(
                            result.loc[result.index[i], upper_col]) if upper_col in result.columns else None

                        result_data.append({
                            'forecast_year': forecast_year,
                            'province': province,
                            'forecast_value': forecast_value,
                            'lower_bound': lower_bound,
                            'upper_bound': upper_bound
                        })

            return {'success': True, 'data': result_data}
        else:
            return {'success': False, 'error': 'MinT预测返回空结果'}

    except Exception as e:
        return {'success': False, 'error': str(e)}


def _execute_ilr_var_forecast(data_df, horizon):
    """执行ILR-VAR预测"""
    try:
        # 先准备数据格式
        from src.data_loader import DataLoader
        
        # 创建数据加载器并处理数据
        data_loader = DataLoader()
        
        # 将DataFrame转换为DataLoader期望的格式
        processed_data = data_loader.process_dataframe(data_df)
        
        # 包装数据以符合Framework3的期望格式
        wrapped_data = {
            'Sheet1': processed_data
        }
        
        # 使用Framework3进行ILR-VAR预测
        framework3 = Framework3(wrapped_data)

        # 执行预测
        result = framework3.forecast(horizon)

        if result is not None:
            # 转换结果格式
            result_data = []
            last_year = data_df.index.max()

            for i in range(horizon):
                forecast_year = last_year + i + 1
                for province in result.columns:
                    if not province.endswith(('_yhat_lower', '_yhat_upper')):
                        forecast_value = float(result.loc[result.index[i], province])
                        lower_col = f'{province}_yhat_lower'
                        upper_col = f'{province}_yhat_upper'

                        lower_bound = float(
                            result.loc[result.index[i], lower_col]) if lower_col in result.columns else None
                        upper_bound = float(
                            result.loc[result.index[i], upper_col]) if upper_col in result.columns else None

                        result_data.append({
                            'forecast_year': forecast_year,
                            'province': province,
                            'forecast_value': forecast_value,
                            'lower_bound': lower_bound,
                            'upper_bound': upper_bound
                        })

            return {'success': True, 'data': result_data}
        else:
            return {'success': False, 'error': 'ILR-VAR预测返回空结果'}

    except Exception as e:
        return {'success': False, 'error': str(e)}


def _save_forecast_results(task, framework, result_data):
    """保存预测结果到数据库"""
    result_objects = []

    for item in result_data:
        result_objects.append(ForecastResult(
            task=task,
            framework=framework,
            forecast_year=item['forecast_year'],
            province=item['province'],
            region=item.get('region', ''),
            forecast_value=item['forecast_value'],
            lower_bound=item.get('lower_bound'),
            upper_bound=item.get('upper_bound')
        ))

    ForecastResult.objects.bulk_create(result_objects, ignore_conflicts=True)


def _safe_float(value):
    """安全地将值转换为float，None值转换为0"""
    return float(value) if value is not None else 0.0
